﻿using System;
using System.Collections.Generic;
using System.Threading;
using Pfz.Caching;

namespace Pfz.Threading
{
	/// <summary>
	/// Base class for DiposableThreadLocal and its generic version.
	/// </summary>
	public class BaseDisposableThreadLocal:
		IDisposable,
		IGarbageCollectionAware
	{
		internal readonly object _lock = new object();
		internal Dictionary<Thread, DisposeAssurer> _dictionary = new Dictionary<Thread, DisposeAssurer>(ReferenceComparer.Instance);

		internal BaseDisposableThreadLocal()
		{
			GCUtils.RegisterForCollectedNotification(this);
		}

		/// <summary>
		/// Releases the values used by all threads that accessed this thread-local variable.
		/// </summary>
		public void Dispose()
		{
			GCUtils.UnregisterFromCollectedNotification(this);

			Dictionary<Thread, DisposeAssurer> dictionary;
			lock(_lock)
			{
				dictionary = _dictionary;

				if (dictionary == null)
					return;

				_dictionary = null;
			}

			foreach(var value in dictionary.Values)
				if (value != null)
					value.Dispose();
		}
		void IGarbageCollectionAware.OnCollected()
		{
			int count = 0;
			KeyValuePair<Thread, DisposeAssurer>[] toRemove;
			lock(_lock)
			{
				var dictionary = _dictionary;
				if (dictionary == null)
				{
					GCUtils.UnregisterFromCollectedNotification(this);
					return;
				}

				try
				{
					toRemove = new KeyValuePair<Thread, DisposeAssurer>[dictionary.Count];
					foreach(var pair in dictionary)
					{
						var thread = pair.Key;
						
						if (!thread.IsAlive)
						{
							toRemove[count] = pair;
							count++;
						}
					}
				}
				catch
				{
					return;
				}

				for(int i=0; i<count; i++)
					dictionary.Remove(toRemove[i].Key);
			}

			for(int i=0; i<count; i++)
				toRemove[i].Value.Dispose();
		}

		internal IDisposable _GetOrCreateValue(Func<IDisposable> valueFactory)
		{
			if (valueFactory == null)
				throw new ArgumentNullException("valueFactory");

			var thread = Thread.CurrentThread;

			if (thread.IsThreadPoolThread)
				throw new InvalidOperationException("Can't create ThreadLocal variables for .Net default ThreadPool. Such threads are revived with all their variables.");

			lock(_lock)
			{
				var dictionary = _dictionary;
				if (dictionary == null)
					throw new ObjectDisposedException(GetType().FullName);

				DisposeAssurer disposeCaller;
				if (_dictionary.TryGetValue(thread, out disposeCaller))
					return disposeCaller.Value;

				var result = valueFactory();
				disposeCaller = DisposeAssurer.Create(result);
				_dictionary.Add(thread, disposeCaller);
				return result;
			}
		}
		internal IDisposable _Value
		{
			get
			{
				var thread = Thread.CurrentThread;
				lock(_lock)
				{
					var dictionary = _dictionary;
					if (dictionary == null)
						throw new ObjectDisposedException(GetType().FullName);

					DisposeAssurer disposeCaller;
					if (_dictionary.TryGetValue(thread, out disposeCaller))
						return disposeCaller.Value;

					return null;
				}
			}
			set
			{
				var thread = Thread.CurrentThread;

				if (thread.IsThreadPoolThread)
					throw new InvalidOperationException("Can't create ThreadLocal variables for .Net default ThreadPool. Such threads are revived with all their variables.");

				DisposeAssurer oldDisposeCaller;

				lock(_lock)
				{
					var dictionary = _dictionary;
					if (dictionary == null)
						throw new ObjectDisposedException(GetType().FullName);

					_dictionary.TryGetValue(thread, out oldDisposeCaller);

					if (value == null)
						_dictionary.Remove(thread);
					else
					{
						DisposeAssurer newDisposeCaller = DisposeAssurer.Create(value);
						_dictionary[thread] = newDisposeCaller;
					}
				}

				if (oldDisposeCaller != null)
					oldDisposeCaller.Dispose();
			}
		}
	}
}
